/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.br.comps;

import coalescent.br.BREvent;
import coalescent.br.BRGenealogy;
import coalescent.br.BRSPopGen;
import coalescent.br.BRWCompInfo;
import coalescent.data.PopGenData;
import coalescent.model.PopGenModel;

import commons.core.Statistic;
import commons.util.FormatUtil;
import commons.util.MapUtil;

import java.math.BigDecimal;
import java.util.*;

/**
 * Builds all the genealogies of a sample configuration by observing its recursion
 * (via post-order implementation).
 *
 * @author Susanta Tewari
 * @version 1.0 Dec 3, 2010
 *
 * @param <E>
 * @param <M>
 * @param <D>
 * @param <S>
 * @param <R>
 */
public class BRWGenealogyBuilder<E extends Enum<E>, M extends PopGenModel<E>,
                                 D extends PopGenData<?, E, M, D>, S extends Statistic<D, S>,
                                 R extends BRSPopGen<?, ?, E, M, D, S, R>>
        extends BRWCompAdapter<E, R> {

    /**
     * cache of genealogies
     */
    private Map<R, List<BRGenealogy<E, M, D, S, R>>> genealogiesCache = new HashMap<>();

    /**
     * the final result of the computation. Holds all the genealogies of the statistic.
     */
    private List<BRGenealogy<E, M, D, S, R>> sampleGenealogies;

    /**
     * Provides the genealogies of the sample configuration (statistic) used in the
     * last computation in postorder traversal of the genealogy tree.
     *
     * @return genealogies of the last sample configuration
     */
    public List<BRGenealogy<E, M, D, S, R>> genealogies() {
        return sampleGenealogies;
    }

    /**
     * Genealogies sorted by their probabilities in the descending order
     *
     * @return genealogies sorted in descending order of their probabilities
     */
    @SuppressWarnings("unchecked")
    public Map<BRGenealogy<E, M, D, S, R>, BigDecimal> genealogiesSortedByProb() {

        Map<BRGenealogy<E, M, D, S, R>, BigDecimal> genealogy_prob_map = new HashMap<>();

        for (BRGenealogy<E, M, D, S, R> genealogy : sampleGenealogies) {

            genealogy_prob_map.put(genealogy, genealogy.probability());
        }


        // sort by prob
        genealogy_prob_map = MapUtil.sortByValue(genealogy_prob_map, Collections.reverseOrder());

        return genealogy_prob_map;
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onBoundaryCondn(BRWCompInfo<E, R> event) {

        R sampleConfig = event.getCurrentConfig();

        if (genealogiesCache.get(sampleConfig) == null) {

            List<BRGenealogy<E, M, D, S, R>> genealogies = new ArrayList<>();

            genealogies.add(new BRGenealogy<>(sampleConfig,
                                              new ArrayList<BREvent<E, M, D, S, R>>()));
            genealogiesCache.put(sampleConfig, genealogies);
        }
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onStartingIteratingEvents(BRWCompInfo<E, R> event) {

        R sampleConfig = event.getCurrentConfig();

        genealogiesCache.put(sampleConfig, new ArrayList<BRGenealogy<E, M, D, S, R>>());
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onIteratingAncestralConfigs(BRWCompInfo<E, R> event) {

        R sampleConfig                               = event.getCurrentConfig();
        List<BRGenealogy<E, M, D, S, R>> genealogies = genealogiesCache.get(sampleConfig);
        R ancestralConfig                            = event.getAncestralConfig();

        for (BRGenealogy<E, M, D, S, R> ancestralGenealogy :
                genealogiesCache.get(ancestralConfig)) {

            List<BREvent<E, M, D, S, R>> ancestralEvents = new ArrayList<>();

            ancestralEvents.add(new BREvent<E, M, D, S, R>(ancestralConfig, event.getEventType()));
            ancestralEvents.addAll(ancestralGenealogy.getEventChain());
            genealogies.add(new BRGenealogy<>(sampleConfig, ancestralEvents));
        }
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onFinishedIteration(BRWCompInfo<E, R> event) {

        super.onFinishedIteration(event);

        R sampleConfig = event.getCurrentConfig();

        if (event.getSampleConfig().equals(sampleConfig)) {

            sampleGenealogies = genealogiesCache.get(sampleConfig);


            // de-allocate the resources
            genealogiesCache = null;
        }
    }

    /**
     * Prints genealogies in descending order of their probabilities. Also computes the sum of those
     * probabilities.
     *
     * @return
     */
    @Override
    public String toString() {

        super.toString();


        // prepare the necessary
        Map<BRGenealogy<E, M, D, S, R>, BigDecimal> genealogy_prob_map = genealogiesSortedByProb();
        BigDecimal totalProb                                           = BigDecimal.valueOf(0.0);


        // start printing
        StringBuilder builder = new StringBuilder();

        builder.append("Genealogies:\n");

        for (BRGenealogy genealogy : sampleGenealogies) {

            BigDecimal prob = genealogy_prob_map.get(genealogy);

            totalProb = totalProb.add(prob);

            builder.append("Event Sequence: " + genealogy.toStringEventSeq() + " Prob: "
                           + FormatUtil.format(prob) + " Genealogy: " + genealogy + "\n");
        }

        builder.append("Sum of genealogy probabilities: " + FormatUtil.format(totalProb) + "\n");

        return builder.toString();
    }
}
